import logging
import util
class HMM:

  def __init__(self, parameters):
    self.p = parameters
    self.p.logAll()
  def run(self, observedSeq):
    self.computeMostLikelyStateSequence(observedSeq)
    self.islands = self.decode(self.trellis)
    return self.islands
    
    
    
  def computeMostLikelyStateSequence(self, observedSeq):
    p = self.p
    
    logger = logging.getLogger("computeMostLikelyStateSequence")
    trellis = []
    
    bt = util.Counter(True)
    delta = util.Counter(True)
    for s in p.stateset:
      if s == "S":
        import math
        delta[s] = math.log(1)
      bt[s] = None
    d = dict()
    d["delta"] = delta
    d["backtrace"] = bt
    trellis.append(d)
    #logger.debug("trellis: "+str(trellis))
    
    for char in observedSeq:    
      bt = util.Counter(True)
      delta = util.Counter(True)
      for s in p.stateset:
        delta[s], bt[s] = self.computeInnerScores(char, s, p, trellis)
      d = dict()
      d["delta"] = delta
      d["backtrace"] = bt
      trellis.append(d)
    logger.debug(trellis)
    self.trellis = trellis
    
  def computeInnerScores(self, char, s, p, trellis):
    m = float("-inf")
    am = None
    #for prevS in trellis[-1].keys():
    for prevS in p.stateset:
      d = p.tokenGiven[s][char] + p.stateGiven[prevS][s] + trellis[-1]["delta"][prevS]
      if m < d:
        m = d;
        am = prevS
    return m, am
    
  def decode(self, trellis):
    mostProbableState = trellis[-1]["delta"].argMax()
    annot = []
    state = "out"
    start = 0
    end = 0
    
    for i in range(len(trellis)- 1, 0, -1):
      if state == "out" and mostProbableState == "I":
        state = "in"
        end = i
      if state == "in" and mostProbableState == "B":
        state = "out"
        start = i + 1
        n = dict()
        n["start"] = start
        n["end"] = end
        n["length"] = end - start + 1
        annot.append(n)
      mostProbableState = trellis[i]["backtrace"][mostProbableState]
    return annot
  def getHiddenSequence(self):
    trellis = self.trellis
    #result = "B"*(len(trellis) - 1)
    result2 = ""
    mostProbableState = trellis[-1]["delta"].argMax()
    print trellis[-1]["delta"].argMax()
    annot = []
    state = "out"
    start = 0
    end = 0
    
    for i in range(len(trellis)- 1, 0, -1):
      if mostProbableState == None:
        break
      result2 = mostProbableState + result2
      if state == "out" and mostProbableState == "I":
        state = "in"
        end = i + 1
      if state == "in" and mostProbableState == "B":
        state = "out"
        start = i + 2
        n = dict()
        n["start"] = start
        n["end"] = end
        n["length"] = end - start + 1
        #result = result[:n['start'] - 1] + 'I'*n['length'] + result[n['end']:]
        annot.append(n)
      mostProbableState = trellis[i]["backtrace"][mostProbableState]
    if state == "in":
      state = "out"
      start = 1
      n = dict()
      n["start"] = start
      n["end"] = end
      n["length"] = end - start + 1
      #result = result[:n['start'] - 1] + 'I'*n['length'] + result[n['end']:]
      annot.append(n)
    #print result
    #print result2
    return result2
  
  def viterbiScore(self):
    return self.trellis[-1]["delta"].max(float("-inf"))




class UnsupervisedHMM:
  def __init__(self, parameters):
    from Statistics import StatisticsCollector
    self.p = parameters
    #self.p = StatisticsCollector(None, None)
  def runHMM(self, testSet):
    p = self.p
    print p
    h = HMM(p)
    scr = float("-inf")
    #self.score(testSet.items()[0][1][0])
    while True:
      for k, v in testSet.items():
        """
        seq = v[0]
        h.computeMostLikelyStateSequence(seq)
        st = h.getHiddenSequence()
        pofo = self.pOfO(seq, st)
        f = self.Foreward(seq)
        b = self.Backward(seq)
        A = dict()
        E = dict()
        for s in p.stateset:
          A[s] = util.Counter(True)
          E[s] = util.Counter(True)
        for i in range(0, len(seq) - 2):	
          for s in p.stateset:
            for prevS in p.stateset:
              A[prevS][s] = self.logsum(A[prevS][s], f[i][prevS] + p.stateGiven[prevS][s] + b[i+1][prevS])
              #A[prevS][s] = self.logsum(A[prevS][s], f[i][prevS] + p.stateGiven[prevS][s] + b[i+1][prevS]+p.tokenGiven[s][seq[i+1]])
            E[s][seq[i]] = self.logsum(E[s][seq[i]], f[i][s] + b[i][s])
        #divide by prob of O
        for s in p.stateset:
          A[s].incrementAll(A[s].keys(), -pofo)
          E[s].incrementAll(E[s].keys(), -pofo)
        #renormalize:
        for s in p.stateset:
          sum = 0.0
          for t in p.stateset:                  
            sum = self.logsum( sum, A[s][t])
          A[s].incrementAll(A[s].keys(), -sum)
        for s in p.stateset:
          sum = 0.0
          for c in E[s]:
            sum = self.logsum( sum, E[s][c])
          E[s].incrementAll(E[s].keys(), -sum)
        print p
        self.p.stateGiven = A
        self.p.tokenGiven = E
        p = self.p
        break
        """
        o = v[0]
        h.computeMostLikelyStateSequence(o)
        st = h.getHiddenSequence()
        alpha = self.Foreward(o)
        beta = self.Backward(o)
        #print alpha
        #print beta
        #import sys
        #sys.exit(1)
        #print str(len(alpha))+", "+str(len(beta))
        #print str(len(o))+", "+str(len(st))
        #print h.trellis
        pofo = self.pOfO(o, st)
        xi = []
        for t in range(len(o) - 1):
          d = dict()
          for i in p.stateset:
            c = util.Counter(True)
            for j in p.stateset:
              c[j] = alpha[t][i] + \
                     p.stateGiven[i][j] + \
                     beta[t+1][j]+ \
                     p.tokenGiven[j][o[t+1]]
            c.incrementAll(c.keys(), -pofo)
            d[i] = c
            
          xi.append(d)
        print len(xi) 
        gamma = []
        for t in range(len(o) - 1):
          c = util.Counter(True)
          for i in p.stateset:
            sum = 0.0
            for j in p.stateset:
              
              sum = self.logsum(sum, xi[t][i][j])
            c[i] = sum
          gamma.append(c)
        pi = gamma[t]
        
        a = dict()
        for i in p.stateset:
          den = 0.0
          for t in range(len(o) - 1):
            den = self.logsum(den, gamma[t][i])
          c = util.Counter(True)
          for j in p.stateset:
            num = 0.0
            for t in range(len(o) - 1):
              #print xi[t][j][i]
              num = self.logsum(num, xi[t][j][i])
            c[j] = num
          c.incrementAll(c.keys(), -den)
          a[i] = c
 
        b = dict()
        for j in p.stateset:
          c = util.Counter(True)
          den = 0.0
          for t in range(len(o) - 1):
            den = self.logsum(den, gamma[t][j])
          for t in range(len(o) - 1):
            c[o[t]] = self.logsum(c[o[t]], gamma[t][j])
          c.incrementAll(c.keys(), -den)
          b[j] = c
        #print xi
        #print gamma 
        #print pi
        #print a
        #print b
        #import sys
        #sys.exit(1)
        print p
        self.p.stateGiven["S"] = pi
        self.p.stateGiven = a
        self.p.tokenGiven = b
        p = self.p

        break
        
      self.score(testSet.items()[0][1][0])
      print h.getHiddenSequence()
      print h.viterbiScore()
      vscr =  h.viterbiScore()
      import math
      if abs(scr - vscr) > 10E-2:
        scr = vscr
      else:
        print h.getHiddenSequence()
        return
      
  def pOfO(self, obs, hid):
    #print len(obs)
    #print len(hid)
    p = self.p
    sum = p.stateGiven["S"][hid[0]]+p.tokenGiven[hid[0]][obs[0]]
    for i in range (1, len(obs)):
      sum += (p.stateGiven[hid[i-1]][hid[i]] \
               +p.tokenGiven[hid[i]][obs[i]])
    return sum
  def score(self, observedSeq):
    p = self.p
    sum = 0.0
    for char in observedSeq:
      for state in p.stateset:
        sum = self.logsum( sum, p.tokenGiven[state][char])
    print sum
    return sum

  def Foreward(self, observedSeq):
    p = self.p
    f = []
    d = util.Counter(True)
    for s in p.stateset:
      if s == "S":
        import math
        d[s] = math.log(1)
    f.append(d)

    for char in observedSeq:
      d = util.Counter(True)
      for s in p.stateset:
        sum = 0.0
        for prevS in p.stateset:
          sum = self.logsum(sum, p.stateGiven[prevS][s] + f[-1][prevS])
        sum += p.tokenGiven[s][char]
        d[s] = sum       
      f.append(d)
    f.pop(0)
    return f
  
  def Backward(self, observedSeq):
    p = self.p
    b = []
    d = util.Counter(True)
    for s in p.stateset:
      import math
      d[s] = math.log(1)
    b.append(d)

    for i in range(len(observedSeq) - 2, -1, -1):
      d = util.Counter(True)
      for s in p.stateset:
        sum = 0.0
        for nextS in p.stateset:
          sum = self.logsum(sum, p.stateGiven[nextS][s] + b[0][nextS] + p.tokenGiven[nextS][observedSeq[i+1]])
        d[s] = sum       
      b.insert(0, d)
    return b
    
  def logsum(self, loga, logb):
    if loga == float("-inf") and logb == float("-inf"):
      return float("-inf")
    import math
    if loga > logb:
      import math
      return loga + math.log(1 + math.exp(logb - loga))
    else:
      if loga == float("-inf") and logb == float("-inf"):
        return float("-inf")
      import math
      return logb + math.log(1 + math.exp(loga - logb))
      
      
